package IntelIpps

/*
#cgo CFLAGS: -ID:/DevelopTools/IPPS2022/ipp/2021.12/include
#cgo LDFLAGS: -LD:/DevelopTools/IPPS2022/ipp/2021.12/lib -lipps -lippcore
#include <ipp.h>
#include <stdlib.h>
*/
import "C"
import (
	"errors"
	"unsafe"
)

type FftState struct {
	Spec   *C.IppsFFTSpec_C_32f
	Buffer []byte
}

var fftStates = make(map[int]*FftState) // len -> state
func (i *IppImpl) FFTInitFloat(length int) error {
	if length <= 0 || (length&(length-1)) != 0 {
		return errors.New("FFT length must be power of 2")
	}

	order := 0
	n := length
	for n > 1 {
		n >>= 1
		order++
	}

	flag := C.int(C.IPP_FFT_NODIV_BY_ANY)

	var sizeSpec, sizeInit, sizeBuffer C.int
	status := C.ippsFFTGetSize_C_32f(C.int(order), flag, C.ippAlgHintNone,
		&sizeSpec, &sizeInit, &sizeBuffer)
	if status != C.ippStsNoErr {
		return errors.New("ippsFFTGetSize_C_32f failed")
	}

	// 使用 ippsMalloc_8u 分配对齐内存
	pSpec := C.ippsMalloc_8u(sizeSpec)
	if pSpec == nil {
		return errors.New("ippsMalloc failed for spec")
	}

	var pSpecBuffer *C.Ipp8u = nil
	if sizeInit > 0 {
		pSpecBuffer = C.ippsMalloc_8u(sizeInit)
		if pSpecBuffer == nil {
			C.ippsFree(unsafe.Pointer(pSpec))
			return errors.New("ippsMalloc failed for init")
		}
	}

	// 工作缓冲区用 Go slice（对齐由 IPP 保证）
	buffer := make([]byte, sizeBuffer)

	// 初始化 FFT
	var spec *C.IppsFFTSpec_C_32f
	status = C.ippsFFTInit_C_32f(&spec, C.int(order), flag, C.ippAlgHintNone,
		pSpec, pSpecBuffer)
	if status != C.ippStsNoErr {
		C.ippsFree(unsafe.Pointer(pSpec))
		if pSpecBuffer != nil {
			C.ippsFree(unsafe.Pointer(pSpecBuffer))
		}
		return errors.New("ippsFFTInit_C_32f failed")
	}

	// 释放临时 init 缓冲区
	if pSpecBuffer != nil {
		C.ippsFree(unsafe.Pointer(pSpecBuffer))
	}

	fftStates[length] = &FftState{
		Spec:   spec,
		Buffer: buffer,
	}

	return nil
}

func (i *IppImpl) FFTFloat(real, imag []float32) []float32 {
	if len(real) != len(imag) {
		return nil
	}
	length := len(real)
	state := fftStates[length]
	if state == nil {
		return nil
	}

	// 复制输入，避免污染
	r := make([]float32, length)
	img := make([]float32, length)
	copy(r, real)
	copy(img, imag)

	// FFT
	C.ippsFFTFwd_CToC_32f(
		(*C.Ipp32f)(unsafe.Pointer(&r[0])),
		(*C.Ipp32f)(unsafe.Pointer(&img[0])),
		(*C.Ipp32f)(unsafe.Pointer(&r[0])),
		(*C.Ipp32f)(unsafe.Pointer(&img[0])),
		state.Spec,
		(*C.Ipp8u)(unsafe.Pointer(&state.Buffer[0])),
	)

	// Magnitude
	out := make([]float32, length)
	C.ippsMagnitude_32f(
		(*C.Ipp32f)(unsafe.Pointer(&r[0])),
		(*C.Ipp32f)(unsafe.Pointer(&img[0])),
		(*C.Ipp32f)(unsafe.Pointer(&out[0])),
		C.int(length),
	)

	//// 频率翻转
	//half := length / 2
	//result := make([]float32, length)
	//copy(result, out[half:])
	//copy(result[half:], out[:half])
	return out
}

func (i *IppImpl) FFTFreeFloat() {
	for _, s := range fftStates {
		if s.Spec != nil {
			C.ippsFree(unsafe.Pointer(s.Spec))
		}
	}
	fftStates = make(map[int]*FftState)
}
